/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <gtest/gtest.h>


#include "graph/utils/tensor_utils.h"
#include "graph/utils/graph_utils.h"
#include "graph/utils/op_desc_utils.h"
#include "graph/utils/attr_utils.h"
#include "graph/debug/ge_attr_define.h"
#define protected public
#define private   public
#include "common/configuration.h"
#include "adapter/common/op_store_adapter_manager.h"
#include "adapter/tbe_adapter/tbe_op_store_adapter.h"
#include "ops_kernel_store/fe_ops_kernel_info_store.h"
#include "ops_store/ops_kernel_manager.h"
#include "graph_optimizer/shape_format_transfer/trans_node_manager/trans_node_manager.h"
#include "graph_optimizer/shape_format_transfer/trans_node_implementation/trans_node_base_generator.h"
#include "graph_optimizer/shape_format_transfer/trans_node_implementation/trans_node_cast_generator.h"
#include "graph_optimizer/shape_format_transfer/trans_node_implementation/trans_node_reshape_generator.h"
#include "graph_optimizer/shape_format_transfer/trans_node_implementation/trans_node_transpose_generator.h"
#include "graph_optimizer/shape_format_transfer/trans_node_implementation/trans_node_transdata_generator.h"
#undef protected
#undef private

#include <iostream>

using namespace std;
using namespace ge;
using namespace fe;

using TbeOpStoreAdapterPtr = std::shared_ptr<TbeOpStoreAdapter>;

class STEST_FE_TRANSOP_INSERT_COMPLEX_2 : public testing::Test {
 protected:
  void SetUp()
  {
    std::map<std::string, std::string> options;
    op_store_adapter_manager_ptr_ = std::make_shared<OpStoreAdapterManager>();
    TbeOpStoreAdapterPtr tbe_adapter_ptr = std::make_shared<TbeOpStoreAdapter>();
    op_store_adapter_manager_ptr_->map_all_op_store_adapter_.emplace(std::make_pair("tbe_op_adapter", tbe_adapter_ptr));
    fe_ops_kernel_info_store_ptr_ = make_shared<fe::FEOpsKernelInfoStore>(op_store_adapter_manager_ptr_);
    FEOpsStoreInfo tbe_custom {
        6,
        "tbe-custom",
        EN_IMPL_HW_TBE,
        "./air/test/engines/nneng/ut/testcase/fusion_engine/ops_kernel_store/fe_config/tbe_custom_opinfo",
        ""};
    vector<FEOpsStoreInfo> store_info;
    store_info.emplace_back(tbe_custom);
    Configuration::Instance(fe::AI_CORE_NAME).ops_store_info_vector_ = (store_info);
    OpsKernelManager::Instance(AI_CORE_NAME).Finalize();

    fe_ops_kernel_info_store_ptr_->Initialize(options);
  }

  void TearDown()
  {
    fe_ops_kernel_info_store_ptr_->Finalize();

  }

  shared_ptr<fe::FEOpsKernelInfoStore> fe_ops_kernel_info_store_ptr_;
  OpStoreAdapterManagerPtr op_store_adapter_manager_ptr_;
 protected:

};


TEST_F(STEST_FE_TRANSOP_INSERT_COMPLEX_2, InsertAllTransop_1) {
  // src:cce op, dst:cce op
  ComputeGraphPtr graph = std::make_shared<ComputeGraph>("test_graph_input");
  OpDescPtr src_op = std::make_shared<OpDesc>("var1", fe::VARIABLE);
  vector<int64_t> dims4_d = {100,200,300,400};
  vector<int64_t> dimsfz = {380000,25,16,16};
  GeTensorDesc src_tensor_desc(GeShape(dims4_d), ge::FORMAT_HWCN, ge::DT_FLOAT16);
  src_tensor_desc.SetOriginShape(GeShape(dims4_d));
  src_tensor_desc.SetOriginFormat(ge::FORMAT_HWCN);
  src_op->AddOutputDesc(src_tensor_desc);
  src_op->AddInputDesc(src_tensor_desc);
  auto src_node = graph->AddNode(src_op);
  ge::AttrUtils::SetInt(src_op, FE_IMPLY_TYPE, -1);
  ge::AttrUtils::SetStr(src_op, PARENT_OP_TYPE, "Variable");

  OpDescPtr trans_op_0 = std::make_shared<OpDesc>("transdata_0", "TransData");
  GeTensorDesc trans_in_tensor_desc(GeShape(dimsfz), ge::FORMAT_FRACTAL_Z, ge::DT_FLOAT16);
  trans_in_tensor_desc.SetOriginShape(GeShape(dims4_d));
  trans_in_tensor_desc.SetOriginFormat(ge::FORMAT_HWCN);
  GeTensorDesc trans_out_tensor_desc(GeShape(dims4_d), ge::FORMAT_HWCN, ge::DT_FLOAT16);
  trans_out_tensor_desc.SetOriginShape(GeShape(dims4_d));
  trans_out_tensor_desc.SetOriginFormat(ge::FORMAT_HWCN);
  trans_op_0->AddInputDesc(trans_in_tensor_desc);
  trans_op_0->AddOutputDesc(trans_out_tensor_desc);
  auto trans_node_0 = graph->AddNode(trans_op_0);
  ge::AttrUtils::SetInt(trans_op_0, FE_IMPLY_TYPE, 6);
  ge::AttrUtils::SetBool(trans_op_0, ge::ATTR_INSERTED_BY_GE, true);
  
  OpDescPtr apply_op = std::make_shared<OpDesc>("apply", "ApplyMomentum");
  GeTensorDesc apply_tensor_desc(GeShape(dimsfz), ge::FORMAT_FRACTAL_Z, ge::DT_FLOAT);
  apply_tensor_desc.SetOriginShape(GeShape(dims4_d));
  apply_tensor_desc.SetOriginFormat(ge::FORMAT_HWCN);
  apply_op->AddInputDesc(apply_tensor_desc);
  apply_op->AddOutputDesc(apply_tensor_desc);
  auto apply_node = graph->AddNode(apply_op);
  ge::AttrUtils::SetInt(apply_op, FE_IMPLY_TYPE, 6);

  OpDescPtr trans_op_1 = std::make_shared<OpDesc>("transdata_1", "TransData");
  trans_op_1->AddInputDesc(trans_out_tensor_desc);
  trans_op_1->AddOutputDesc(trans_in_tensor_desc);
  auto trans_node_1 = graph->AddNode(trans_op_1);
  ge::AttrUtils::SetInt(trans_op_1, FE_IMPLY_TYPE, 6);
  ge::AttrUtils::SetBool(trans_op_1, ge::ATTR_INSERTED_BY_GE, true);
  
  OpDescPtr dst_op = std::make_shared<OpDesc>("var2", fe::VARIABLE);
  GeTensorDesc dst_tensor_desc(GeShape(dims4_d), ge::FORMAT_HWCN, ge::DT_FLOAT16);
  dst_tensor_desc.SetOriginShape(GeShape(dims4_d));
  dst_tensor_desc.SetOriginFormat(ge::FORMAT_HWCN);
  dst_op->AddOutputDesc(dst_tensor_desc);
  dst_op->AddInputDesc(dst_tensor_desc);
  auto dst_node = graph->AddNode(dst_op);
  ge::AttrUtils::SetInt(dst_op, FE_IMPLY_TYPE, -1);
  ge::AttrUtils::SetStr(dst_op, PARENT_OP_TYPE, "Variable");

  GraphUtils::AddEdge(src_node->GetOutDataAnchor(0), trans_node_0->GetInDataAnchor(0));
  GraphUtils::AddEdge(trans_node_0->GetOutDataAnchor(0), apply_node->GetInDataAnchor(0));
  GraphUtils::AddEdge(apply_node->GetOutDataAnchor(0), trans_node_1->GetInDataAnchor(0));
  GraphUtils::AddEdge(trans_node_1->GetOutDataAnchor(0), dst_node->GetInDataAnchor(0));

  TransNodeManager trans_op_insert(fe_ops_kernel_info_store_ptr_);
  trans_op_insert.Initialize();
  Status status = trans_op_insert.InsertAndMergeTransNodes(*(graph.get()));

  int count_node = 0;
  uint32_t index = 0;
  ASSERT_EQ(fe::SUCCESS, status);
  for (auto node : graph->GetDirectNode()) {
    ASSERT_NE(node, nullptr);
    count_node++;
    EXPECT_NE(node->GetName(),"transdata_0");
    EXPECT_NE(node->GetName(),"transdata_1");
    if (node->GetName() == "apply") {
      {
        ge::GeShape shape = node->GetOpDesc()->GetInputDescPtr(0)->GetShape();
        EXPECT_EQ(shape.GetDimNum(), 4);
        EXPECT_EQ(shape.GetDims(), dimsfz);
        EXPECT_EQ(ge::FORMAT_FRACTAL_Z,
                  node->GetOpDesc()->GetInputDescPtr(0)->GetFormat());
        EXPECT_EQ(ge::DT_FLOAT,
                  node->GetOpDesc()->GetInputDescPtr(0)->GetDataType());
      }
      {
        ge::GeShape shape = node->GetOpDesc()->GetOutputDescPtr(0)->GetShape();
        EXPECT_EQ(shape.GetDimNum(), 4);
        EXPECT_EQ(shape.GetDims(), dimsfz);
        EXPECT_EQ(ge::FORMAT_FRACTAL_Z,
                  node->GetOpDesc()->GetOutputDescPtr(0)->GetFormat());
        EXPECT_EQ(ge::DT_FLOAT,
                  node->GetOpDesc()->GetOutputDescPtr(0)->GetDataType());
      }
    }
    if (node->GetType() == "Cast") {
      if (index == 0) {
        {
          ge::GeShape shape = node->GetOpDesc()->GetInputDescPtr(0)->GetShape();
          EXPECT_EQ(shape.GetDimNum(), 4);
          EXPECT_EQ(shape.GetDims(), dimsfz);
          EXPECT_EQ(ge::FORMAT_FRACTAL_Z,
                    node->GetOpDesc()->GetInputDescPtr(0)->GetFormat());
          EXPECT_EQ(ge::DT_FLOAT16,
                    node->GetOpDesc()->GetInputDescPtr(0)->GetDataType());
        }
        {
          ge::GeShape shape = node->GetOpDesc()->GetOutputDescPtr(0)->GetShape();
          EXPECT_EQ(shape.GetDimNum(), 4);
          EXPECT_EQ(shape.GetDims(), dimsfz);
          EXPECT_EQ(ge::FORMAT_FRACTAL_Z,
                    node->GetOpDesc()->GetOutputDescPtr(0)->GetFormat());
          EXPECT_EQ(ge::DT_FLOAT,
                    node->GetOpDesc()->GetOutputDescPtr(0)->GetDataType());
        }
        index++;
      } else {
        {
          ge::GeShape shape = node->GetOpDesc()->GetInputDescPtr(0)->GetShape();
          EXPECT_EQ(shape.GetDimNum(), 4);
          EXPECT_EQ(shape.GetDims(), dimsfz);
          EXPECT_EQ(ge::FORMAT_FRACTAL_Z,
                    node->GetOpDesc()->GetInputDescPtr(0)->GetFormat());
          EXPECT_EQ(ge::DT_FLOAT,
                    node->GetOpDesc()->GetInputDescPtr(0)->GetDataType());
        }
        {
          ge::GeShape shape = node->GetOpDesc()->GetOutputDescPtr(0)->GetShape();
          EXPECT_EQ(shape.GetDimNum(), 4);
          EXPECT_EQ(shape.GetDims(), dimsfz);
          EXPECT_EQ(ge::FORMAT_FRACTAL_Z,
                    node->GetOpDesc()->GetOutputDescPtr(0)->GetFormat());
          EXPECT_EQ(ge::DT_FLOAT16,
                    node->GetOpDesc()->GetOutputDescPtr(0)->GetDataType());
        }
      }
    }
  }
  EXPECT_EQ(count_node, 7);
}



TEST_F(STEST_FE_TRANSOP_INSERT_COMPLEX_2, InsertAllTransop_2) {
  // src:cce op, dst:cce op
  ComputeGraphPtr graph = std::make_shared<ComputeGraph>("test_graph_input");
  OpDescPtr src_op = std::make_shared<OpDesc>("var1", OP_TYPE_PLACE_HOLDER);
  vector<int64_t> dims4_d = {100,200,300,400};
  vector<int64_t> dimsfz = {380000,25,16,16};
  GeTensorDesc src_tensor_desc(GeShape(dims4_d), ge::FORMAT_HWCN, ge::DT_FLOAT16);
  src_tensor_desc.SetOriginShape(GeShape(dims4_d));
  src_tensor_desc.SetOriginFormat(ge::FORMAT_HWCN);
  src_op->AddOutputDesc(src_tensor_desc);
  src_op->AddInputDesc(src_tensor_desc);
  auto src_node = graph->AddNode(src_op);
  ge::AttrUtils::SetInt(src_op, FE_IMPLY_TYPE, -1);
  ge::AttrUtils::SetStr(src_op, PARENT_OP_TYPE, "Variable");

  OpDescPtr trans_op_0 = std::make_shared<OpDesc>("transdata_0", "TransData");
  GeTensorDesc trans_in_tensor_desc(GeShape(dimsfz), ge::FORMAT_FRACTAL_Z, ge::DT_FLOAT16);
  trans_in_tensor_desc.SetOriginShape(GeShape(dims4_d));
  trans_in_tensor_desc.SetOriginFormat(ge::FORMAT_HWCN);
  trans_in_tensor_desc.SetOriginDataType(ge::DT_FLOAT16);
  GeTensorDesc trans_out_tensor_desc(GeShape(dims4_d), ge::FORMAT_HWCN, ge::DT_FLOAT16);
  trans_out_tensor_desc.SetOriginShape(GeShape(dims4_d));
  trans_out_tensor_desc.SetOriginFormat(ge::FORMAT_HWCN);
  trans_out_tensor_desc.SetOriginDataType(ge::DT_FLOAT16);
  trans_op_0->AddInputDesc(trans_in_tensor_desc);
  trans_op_0->AddOutputDesc(trans_out_tensor_desc);
  auto trans_node_0 = graph->AddNode(trans_op_0);
  ge::AttrUtils::SetInt(trans_op_0, FE_IMPLY_TYPE, 6);

  OpDescPtr apply_op = std::make_shared<OpDesc>("apply", "ApplyMomentum");
  GeTensorDesc apply_tensor_desc(GeShape(dimsfz), ge::FORMAT_FRACTAL_Z, ge::DT_FLOAT);
  apply_tensor_desc.SetOriginShape(GeShape(dims4_d));
  apply_tensor_desc.SetOriginFormat(ge::FORMAT_HWCN);
  apply_op->AddInputDesc(apply_tensor_desc);
  apply_op->AddOutputDesc(apply_tensor_desc);
  auto apply_node = graph->AddNode(apply_op);
  ge::AttrUtils::SetInt(apply_op, FE_IMPLY_TYPE, 6);

  OpDescPtr trans_op_1 = std::make_shared<OpDesc>("transdata_1", "TransData");
  trans_op_1->AddInputDesc(trans_out_tensor_desc);
  trans_op_1->AddOutputDesc(trans_in_tensor_desc);
  auto trans_node_1 = graph->AddNode(trans_op_1);
  ge::AttrUtils::SetInt(trans_op_1, FE_IMPLY_TYPE, 6);

  OpDescPtr dst_op = std::make_shared<OpDesc>("var2", OP_TYPE_END);
  GeTensorDesc dst_tensor_desc(GeShape(dims4_d), ge::FORMAT_HWCN, ge::DT_FLOAT16);
  dst_tensor_desc.SetOriginShape(GeShape(dims4_d));
  dst_tensor_desc.SetOriginFormat(ge::FORMAT_HWCN);
  dst_op->AddOutputDesc(dst_tensor_desc);
  dst_op->AddInputDesc(dst_tensor_desc);
  auto dst_node = graph->AddNode(dst_op);
  ge::AttrUtils::SetInt(dst_op, FE_IMPLY_TYPE, -1);
  ge::AttrUtils::SetStr(dst_op, PARENT_OP_TYPE, "Variable");

  GraphUtils::AddEdge(src_node->GetOutDataAnchor(0), trans_node_0->GetInDataAnchor(0));
  GraphUtils::AddEdge(trans_node_0->GetOutDataAnchor(0), apply_node->GetInDataAnchor(0));
  GraphUtils::AddEdge(apply_node->GetOutDataAnchor(0), trans_node_1->GetInDataAnchor(0));
  GraphUtils::AddEdge(trans_node_1->GetOutDataAnchor(0), dst_node->GetInDataAnchor(0));

  TransNodeManager trans_op_insert(fe_ops_kernel_info_store_ptr_);
  trans_op_insert.Initialize();
  Status status = trans_op_insert.InsertAndMergeTransNodes(*(graph.get()));

  int count_node = 0;
  uint32_t index_cast = 0;
  uint32_t index_transdata = 0;
  ASSERT_EQ(fe::SUCCESS, status);
  for (auto node : graph->GetDirectNode()) {
    ASSERT_NE(node, nullptr);
    count_node++;
    EXPECT_NE(node->GetName(),"transdata_0");
    EXPECT_NE(node->GetName(),"transdata_1");
    if (node->GetName() == "apply") {
      {
        ge::GeShape shape = node->GetOpDesc()->GetInputDescPtr(0)->GetShape();
        EXPECT_EQ(shape.GetDimNum(), 4);
        EXPECT_EQ(shape.GetDims(), dimsfz);
        EXPECT_EQ(ge::FORMAT_FRACTAL_Z,
                  node->GetOpDesc()->GetInputDescPtr(0)->GetFormat());
        EXPECT_EQ(ge::DT_FLOAT,
                  node->GetOpDesc()->GetInputDescPtr(0)->GetDataType());
      }
      {
        ge::GeShape shape = node->GetOpDesc()->GetOutputDescPtr(0)->GetShape();
        EXPECT_EQ(shape.GetDimNum(), 4);
        EXPECT_EQ(shape.GetDims(), dimsfz);
        EXPECT_EQ(ge::FORMAT_FRACTAL_Z,
                  node->GetOpDesc()->GetOutputDescPtr(0)->GetFormat());
        EXPECT_EQ(ge::DT_FLOAT,
                  node->GetOpDesc()->GetOutputDescPtr(0)->GetDataType());
      }
    }
    if (node->GetType() == "Cast") {
      if (index_cast == 0) {
        {
          ge::GeShape shape = node->GetOpDesc()->GetInputDescPtr(0)->GetShape();
          EXPECT_EQ(shape.GetDimNum(), 4);
          EXPECT_EQ(shape.GetDims(), dimsfz);
          EXPECT_EQ(ge::FORMAT_FRACTAL_Z,
                    node->GetOpDesc()->GetInputDescPtr(0)->GetFormat());
          EXPECT_EQ(ge::DT_FLOAT16,
                    node->GetOpDesc()->GetInputDescPtr(0)->GetDataType());
        }
        {
          ge::GeShape shape = node->GetOpDesc()->GetOutputDescPtr(0)->GetShape();
          EXPECT_EQ(shape.GetDimNum(), 4);
          EXPECT_EQ(shape.GetDims(), dimsfz);
          EXPECT_EQ(ge::FORMAT_FRACTAL_Z,
                    node->GetOpDesc()->GetOutputDescPtr(0)->GetFormat());
          EXPECT_EQ(ge::DT_FLOAT,
                    node->GetOpDesc()->GetOutputDescPtr(0)->GetDataType());
        }
        index_cast++;
      } else {
        {
          ge::GeShape shape = node->GetOpDesc()->GetInputDescPtr(0)->GetShape();
          EXPECT_EQ(shape.GetDimNum(), 4);
          EXPECT_EQ(shape.GetDims(), dimsfz);
          EXPECT_EQ(ge::FORMAT_FRACTAL_Z,
                    node->GetOpDesc()->GetInputDescPtr(0)->GetFormat());
          EXPECT_EQ(ge::DT_FLOAT,
                    node->GetOpDesc()->GetInputDescPtr(0)->GetDataType());
        }
        {
          ge::GeShape shape = node->GetOpDesc()->GetOutputDescPtr(0)->GetShape();
          EXPECT_EQ(shape.GetDimNum(), 4);
          EXPECT_EQ(shape.GetDims(), dimsfz);
          EXPECT_EQ(ge::FORMAT_FRACTAL_Z,
                    node->GetOpDesc()->GetOutputDescPtr(0)->GetFormat());
          EXPECT_EQ(ge::DT_FLOAT16,
                    node->GetOpDesc()->GetOutputDescPtr(0)->GetDataType());
        }
      }
    }
    if (node->GetType() == "TransData") {
      if (index_transdata == 0) {
        {
          ge::GeShape shape = node->GetOpDesc()->GetInputDescPtr(0)->GetShape();
          EXPECT_EQ(shape.GetDimNum(), 4);
          EXPECT_EQ(shape.GetDims(), dims4_d);
          EXPECT_EQ(ge::FORMAT_HWCN,
                    node->GetOpDesc()->GetInputDescPtr(0)->GetFormat());
          EXPECT_EQ(ge::DT_FLOAT16,
                    node->GetOpDesc()->GetInputDescPtr(0)->GetDataType());
        }
        {
          ge::GeShape shape = node->GetOpDesc()->GetOutputDescPtr(0)->GetShape();
          EXPECT_EQ(shape.GetDimNum(), 4);
          EXPECT_EQ(shape.GetDims(), dimsfz);
          EXPECT_EQ(ge::FORMAT_FRACTAL_Z,
                    node->GetOpDesc()->GetOutputDescPtr(0)->GetFormat());
          EXPECT_EQ(ge::DT_FLOAT16,
                    node->GetOpDesc()->GetOutputDescPtr(0)->GetDataType());
        }
        index_transdata++;
      } else {
        {
          ge::GeShape shape = node->GetOpDesc()->GetInputDescPtr(0)->GetShape();
          EXPECT_EQ(shape.GetDimNum(), 4);
          EXPECT_EQ(shape.GetDims(), dimsfz);
          EXPECT_EQ(ge::FORMAT_FRACTAL_Z,
                    node->GetOpDesc()->GetInputDescPtr(0)->GetFormat());
          EXPECT_EQ(ge::DT_FLOAT16,
                    node->GetOpDesc()->GetInputDescPtr(0)->GetDataType());
        }
        {
          ge::GeShape shape = node->GetOpDesc()->GetOutputDescPtr(0)->GetShape();
          EXPECT_EQ(shape.GetDimNum(), 4);
          EXPECT_EQ(shape.GetDims(), dims4_d);
          EXPECT_EQ(ge::FORMAT_HWCN,
                    node->GetOpDesc()->GetOutputDescPtr(0)->GetFormat());
          EXPECT_EQ(ge::DT_FLOAT16,
                    node->GetOpDesc()->GetOutputDescPtr(0)->GetDataType());
        }
      }
    }
  }
  EXPECT_EQ(count_node, 7);
}



TEST_F(STEST_FE_TRANSOP_INSERT_COMPLEX_2, InsertPermuteNode)
{
  // src:cce op, dst:cce op
  ComputeGraphPtr graph = std::make_shared<ComputeGraph>("test_graph");
  std::vector<int64_t> dim_src = {1, 1024, 256, 512};
  OpDescPtr src_op = std::make_shared<OpDesc>("B", "B");
  GeTensorDesc src_tensor_desc(GeShape(dim_src), ge::FORMAT_NCHW, ge::DT_FLOAT);
  src_tensor_desc.SetOriginShape(GeShape(dim_src));
  src_tensor_desc.SetOriginFormat(ge::FORMAT_NCHW);
  src_op->AddOutputDesc(src_tensor_desc);
  src_op->AddInputDesc(src_tensor_desc);
  auto src_node = graph->AddNode(src_op);
  ge::AttrUtils::SetInt(src_op, FE_IMPLY_TYPE, 6);

  std::vector<int64_t> dim_dst = {1, 3, 4, 2};
  OpDescPtr dst_op = std::make_shared<OpDesc>("D", "D");
  GeTensorDesc dst_tensor_desc(GeShape(dim_dst), ge::FORMAT_NHWC, ge::DT_FLOAT);
  dst_tensor_desc.SetOriginShape(GeShape(dim_dst));
  dst_tensor_desc.SetOriginFormat(ge::FORMAT_NHWC);
  dst_op->AddInputDesc(dst_tensor_desc);
  dst_op->AddOutputDesc(dst_tensor_desc);
  auto dst_node = graph->AddNode(dst_op);
  ge::AttrUtils::SetInt(dst_op, FE_IMPLY_TYPE, 6);
  vector<bool> input_const_vector = {false};
  dst_op->SetIsInputConst(input_const_vector);

  GraphUtils::AddEdge(src_node->GetOutDataAnchor(0), dst_node->GetInDataAnchor(0));

  TransNodeManager trans_op_insert(fe_ops_kernel_info_store_ptr_);
  trans_op_insert.Initialize();
  Status status = trans_op_insert.InsertAndMergeTransNodes(*(graph.get()));

  ASSERT_EQ(fe::SUCCESS, status);
  int count_node = 0;
  for (auto node : graph->GetDirectNode()) {
    ASSERT_NE(node, nullptr);
    count_node++;
    EXPECT_NE(node->GetType(), "TransposeD");
    if (node->GetType() == "D") {
      {
        ge::GeShape shape_check = node->GetOpDesc()->GetInputDescPtr(0)->GetShape();
        auto size = shape_check.GetDimNum();
        ASSERT_EQ(size, 4);
        vector<int64_t> input_vec_of_b =shape_check.GetDims();
        EXPECT_EQ(input_vec_of_b[0], 1);
        EXPECT_EQ(input_vec_of_b[1], 3);
        EXPECT_EQ(input_vec_of_b[2], 4);
        EXPECT_EQ(input_vec_of_b[3], 2);
      }
      {
        ge::GeShape shape_check = node->GetOpDesc()->GetOutputDescPtr(0)->GetShape();
        auto size = shape_check.GetDimNum();
        ASSERT_EQ(size, 4);
        vector<int64_t> input_vec_of_b =shape_check.GetDims();
        EXPECT_EQ(input_vec_of_b[0], 1);
        EXPECT_EQ(input_vec_of_b[1], 3);
        EXPECT_EQ(input_vec_of_b[2], 4);
        EXPECT_EQ(input_vec_of_b[3], 2);
      }

    }

  }
  EXPECT_EQ(count_node, 2);
}
